import numpy as np
from scipy.sparse import csr_matrix





def compliance_tensor_matrix(tspace,mu):
    E = 2.0*mu
    ldof = tspace.number_of_local_dofs()
    tdim = tspace.tensor_dimension()
    gdim = tspace.geo_dimension()
    bcs, ws = tspace.integrator.quadpts, tspace.integrator.weights
    rho = tspace.mesh.bc_to_point(bcs)[...,0] #(NQ,NC)


    NC = tspace.mesh.number_of_cells()
    NQ = bcs.shape[0]

    phi = tspace.basis(bcs).reshape(NQ,NC,-1,tdim)#(NQ,NC,ldof,tdim)

    #construct matrix
    d = np.array([1, 1, 2])
    M = np.einsum('i,ij, ijkm, m, ijom, j->jko', ws, 1.0/rho, phi/E, d, phi, tspace.mesh.entity_measure(), optimize=True)

    I = np.einsum('ij, k->ijk', tspace.cell2dof.reshape(NC,-1), np.ones(ldof))
    J = I.swapaxes(-1, -2)
    tgdof = tspace.number_of_global_dofs()

    M = csr_matrix((M.flat, (I.flat, J.flat)), shape=(tgdof, tgdof))

    return M
